# -*- coding: utf-8 -*-
"""
缓存管理工具模块

此模块提供用于管理数据缓存的函数。
主要功能包括：
- 缓存数据到文件
- 从缓存文件加载数据
- 清理缓存
"""

import os
import numpy as np
import shutil
from datetime import datetime

from dry_wet_abrupt.config import CACHE_DIR


def cache_data(data, cache_file):
    """
    将数据缓存到文件
    
    参数:
        data: 要缓存的数据
        cache_file (str): 缓存文件路径
        
    返回:
        bool: 是否成功缓存
    """
    try:
        # 确保缓存目录存在
        os.makedirs(os.path.dirname(cache_file), exist_ok=True)
        
        # 保存数据到文件
        np.save(cache_file, data)
        
        print(f"数据已缓存到: {cache_file}")
        return True
    except Exception as e:
        print(f"缓存数据失败: {e}")
        return False


def load_cached_data(cache_file):
    """
    从缓存文件加载数据
    
    参数:
        cache_file (str): 缓存文件路径
        
    返回:
        缓存的数据，如果加载失败则返回None
    """
    try:
        # 检查缓存文件是否存在
        if not os.path.exists(cache_file):
            print(f"缓存文件不存在: {cache_file}")
            return None
        
        # 加载数据
        data = np.load(cache_file, allow_pickle=True).item()
        
        print(f"已从缓存加载数据: {cache_file}")
        return data
    except Exception as e:
        print(f"加载缓存数据失败: {e}")
        return None


def clear_cache(cache_dir=None, older_than=None):
    """
    清理缓存
    
    参数:
        cache_dir (str): 缓存目录，如为None则使用默认缓存目录
        older_than (int): 清理指定天数之前的缓存，如为None则清理所有缓存
        
    返回:
        int: 清理的文件数量
    """
    # 如果未指定缓存目录，使用默认缓存目录
    if cache_dir is None:
        cache_dir = CACHE_DIR
    
    # 如果缓存目录不存在，则无需清理
    if not os.path.exists(cache_dir):
        print(f"缓存目录不存在: {cache_dir}")
        return 0
    
    # 获取当前时间
    now = datetime.now()
    
    # 清理文件计数
    count = 0
    
    # 遍历缓存目录中的所有文件
    for file_name in os.listdir(cache_dir):
        file_path = os.path.join(cache_dir, file_name)
        
        # 如果是文件
        if os.path.isfile(file_path):
            # 如果指定了清理指定天数之前的缓存
            if older_than is not None:
                # 获取文件修改时间
                mtime = datetime.fromtimestamp(os.path.getmtime(file_path))
                
                # 计算文件修改时间距今的天数
                days = (now - mtime).days
                
                # 如果文件修改时间在指定天数之前，则删除
                if days >= older_than:
                    os.remove(file_path)
                    count += 1
            else:
                # 否则删除所有文件
                os.remove(file_path)
                count += 1
    
    print(f"已清理{count}个缓存文件")
    return count


def get_cache_size(cache_dir=None):
    """
    获取缓存大小
    
    参数:
        cache_dir (str): 缓存目录，如为None则使用默认缓存目录
        
    返回:
        int: 缓存大小（字节）
    """
    # 如果未指定缓存目录，使用默认缓存目录
    if cache_dir is None:
        cache_dir = CACHE_DIR
    
    # 如果缓存目录不存在，则返回0
    if not os.path.exists(cache_dir):
        return 0
    
    # 计算缓存目录大小
    total_size = 0
    for dirpath, dirnames, filenames in os.walk(cache_dir):
        for f in filenames:
            fp = os.path.join(dirpath, f)
            total_size += os.path.getsize(fp)
    
    return total_size


def format_size(size_bytes):
    """
    格式化文件大小
    
    参数:
        size_bytes (int): 文件大小（字节）
        
    返回:
        str: 格式化后的文件大小
    """
    # 定义单位
    units = ['B', 'KB', 'MB', 'GB', 'TB']
    
    # 计算单位索引
    unit_index = 0
    while size_bytes >= 1024 and unit_index < len(units) - 1:
        size_bytes /= 1024
        unit_index += 1
    
    # 格式化大小
    return f"{size_bytes:.2f} {units[unit_index]}"